#!/usr/bin/env python3
"""Script for resuming training from command line."""
# pylint: disable=no-name-in-module
import click

from garage import wrap_experiment  # pylint: disable=import-self
from garage.trainer import TFTrainer


@click.group()
def cli():  # noqa: D103
    """The main command group."""


@cli.command()
@click.argument('from_dir')
@click.option(
    '--log_dir',
    default=None,
    help='Log path for resumed experiment. If not specified, will be the same '
    'as from_dir.')
# pylint: disable=bad-docstring-quotes
@click.option('--from_epoch',
              default='last',
              help='When there are multiple snapshots, '
              'specify the index of epoch to restore from. '
              'Can be "first", "last" or a number. '
              'Not applicable when snapshot_mode="last"')
def resume(from_dir, from_epoch, log_dir):
    # pylint: disable=missing-param-doc, missing-type-doc
    """Resume from experiment saved in FROM_DIR."""
    if log_dir is None:
        log_dir = from_dir

    @wrap_experiment(log_dir=log_dir)
    def run_resume(ctxt=None):
        with TFTrainer(snapshot_config=ctxt) as runner:
            runner.restore(from_dir=from_dir, from_epoch=from_epoch)
            runner.resume()

    run_resume()


if __name__ == '__main__':
    cli()
